# geotiff转换工具
import os
from src.config.config import ROOT_DIR, OUTPUT_DIR, TEST_SAVE_DIR
from osgeo import gdal

gdal.UseExceptions()


# os.environ['PROJ_LIB'] = r'D:\commandFiles\anaconda3\envs\paddleseg\Library\share\proj'

def add_geotiff_metadata(src_tif, output_png, dst_geotif):
    # 读取原始TIF的地理信息
    src_ds = gdal.Open(src_tif)
    geotrans = src_ds.GetGeoTransform()
    proj = src_ds.GetProjection()

    # 将预测结果PNG转换为带地理信息的GeoTIFF
    png_ds = gdal.Open(output_png)
    driver = gdal.GetDriverByName("GTiff")
    dst_ds = driver.CreateCopy(dst_geotif, png_ds)
    dst_ds.SetGeoTransform(geotrans)
    dst_ds.SetProjection(proj)
    dst_ds.FlushCache()


if __name__ == '__main__':
    tif_dir = os.path.join(os.path.join(ROOT_DIR, '..', 'data1', 'images'))
    pre_label_dir = os.path.join(os.path.join(TEST_SAVE_DIR, 'pseudo_color_prediction'))
    output_dir = os.path.join(os.path.join(TEST_SAVE_DIR, 'prediction_geotiff'))

    for name in os.listdir(tif_dir):
        if name.endswith('.tif'):
            tif_path = os.path.join(tif_dir, name)
            pre_label_path = os.path.join(pre_label_dir, name.split('.')[0] + '.png')
            output_path = os.path.join(output_dir, name)
            add_geotiff_metadata(tif_path, pre_label_path, output_path)
            print(output_path)
